Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes.

Affected area

• FE OSS kernels or CuTeDSL

Summary

Fix the varlen indexer top-k kernel so that out-of-bounds lanes from predicated vector loads do not participate in radix histogram construction or candidate collection.

The change applies the existing vector-lane bounds predicate during:

• coarse histogram construction;
• direct candidate collection;
• initial refinement histogram and candidate spill.

A regression test covers rows where the valid input prefix contains fewer finite values than top_k, making the top-k threshold -inf.

Why

The kernel uses predicated vector loads and fills out-of-bounds register lanes with -inf. However, those lanes were still counted in the radix histogram and could subsequently be collected as top-k candidates.

When the real top-k threshold was also -inf, the kernel could not distinguish valid in-range -inf values from the synthetic OOB values. This produced thousands of invalid candidates, overflowed the per-row extra buffer, and resulted in an illegal memory access.

Filtering candidates by their logical input index fixes the issue while preserving valid top-k semantics for real in-range -inf values.

Related issues

None.

API and compatibility impact

None. There are no public API or input-contract changes.

The kernel now excludes internal OOB vector lanes from top-k selection. Valid in-range values, including -inf, retain their existing top-k semantics.

Testing

Tested on an NVIDIA B200 (SM100) with CUDA launch blocking enabled.

• python -m pytest -q --tb=short fe_api/dsa/test_DSA_indexer_top_k.py
  • 5 passed
• Replayed the original failing input:
  • shape: (633, 768)
  • dtype: float32
  • seq_lens=633
  • top_k=512
  • completed without IMA
• Replayed the original failing input for 100 consecutive iterations:
  • all iterations completed without IMA
• git diff --check
  • passed

Summary by CodeRabbit

• Bug Fixes

  • Corrected top-k selection to ignore padded vector lanes during histogram processing and index collection.
  • Prevented out-of-range padded values from influencing top-k results, including inputs containing negative infinity.

• Tests

  • Added coverage verifying returned indices remain within valid sequence bounds and match expected top-k values.